doctra 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

doctra/cli/main.py

@@ -9,6 +9,7 @@ detection results, and analyze document structure from the command line.
 import click
 import os
 import sys
+import traceback
 from pathlib import Path
 from typing import Optional
 
@@ -25,6 +26,10 @@ except ImportError:
     from doctra.parsers.enhanced_pdf_parser import EnhancedPDFParser
     from doctra.parsers.table_chart_extractor import ChartTablePDFParser
 
+# Import additional modules
+from doctra.engines.layout.paddle_layout import PaddleLayoutEngine
+from doctra.engines.image_restoration import DocResEngine
+
 
 @click.group(invoke_without_command=True)
 @click.pass_context
@@ -247,7 +252,6 @@ def parse(pdf_path: Path, output_dir: Optional[Path], use_vlm: bool,
     except Exception as e:
         click.echo(f"❌ Error initializing parser: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
 
@@ -271,7 +275,6 @@ def parse(pdf_path: Path, output_dir: Optional[Path], use_vlm: bool,
     except Exception as e:
         click.echo(f"❌ Error during parsing: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
     finally:
@@ -394,7 +397,6 @@ def enhance(pdf_path: Path, output_dir: Optional[Path], restoration_task: str,
     except Exception as e:
         click.echo(f"❌ Error initializing enhanced parser: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
 
@@ -418,7 +420,6 @@ def enhance(pdf_path: Path, output_dir: Optional[Path], restoration_task: str,
     except Exception as e:
         click.echo(f"❌ Error during enhanced parsing: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
     finally:
@@ -526,7 +527,6 @@ def charts(pdf_path: Path, output_dir: Path, use_vlm: bool, vlm_provider: str,
     except Exception as e:
         click.echo(f"❌ Error during chart extraction: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
 
@@ -604,7 +604,6 @@ def tables(pdf_path: Path, output_dir: Path, use_vlm: bool, vlm_provider: str,
     except Exception as e:
         click.echo(f"❌ Error during table extraction: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
 
@@ -683,7 +682,6 @@ def both(pdf_path: Path, output_dir: Path, use_vlm: bool, vlm_provider: str,
     except Exception as e:
         click.echo(f"❌ Error during extraction: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
 
@@ -772,7 +770,6 @@ def visualize(pdf_path: Path, pages: int, columns: int, width: int,
     except Exception as e:
         click.echo(f"❌ Error creating visualization: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
 
@@ -805,7 +802,6 @@ def analyze(pdf_path: Path, dpi: int, min_score: float, layout_model: str, verbo
         click.echo(f"🔍 Analyzing: {pdf_path.name}")
 
         # Create layout engine for analysis only
-        from doctra.engines.layout.paddle_layout import PaddleLayoutEngine
 
         if verbose:
             click.echo(f"   Using model: {layout_model}")
@@ -903,7 +899,6 @@ def analyze(pdf_path: Path, dpi: int, min_score: float, layout_model: str, verbo
     except Exception as e:
         click.echo(f"❌ Error analyzing PDF: {e}", err=True)
         if verbose:
-            import traceback
             click.echo(traceback.format_exc(), err=True)
         sys.exit(1)
 
@@ -922,7 +917,6 @@ def info():
     click.echo("=" * 50)
 
     # Check Python version
-    import sys
     python_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
     click.echo(f"Python version: {python_version}")
 
@@ -1003,7 +997,6 @@ def info():
     # DocRes information
     click.echo("\nDocRes Image Restoration:")
     try:
-        from doctra.engines.image_restoration import DocResEngine
         docres = DocResEngine()
         click.echo(f"  ✅ DocRes available - {len(docres.get_supported_tasks())} restoration tasks")
         click.echo("  Tasks: dewarping, deshadowing, appearance, deblurring, binarization, end2end")

doctra/cli/utils.py

@@ -7,8 +7,10 @@ different CLI commands.
 
 import click
 import sys
+import traceback
 from typing import Optional, Dict, Any
 from pathlib import Path
+from doctra.utils.progress import create_beautiful_progress_bar, create_notebook_friendly_bar
 
 
 def validate_vlm_config(use_vlm: bool, vlm_api_key: Optional[str]) -> None:
@@ -58,7 +60,6 @@ def handle_exception(e: Exception, verbose: bool = False) -> None:
     """
     click.echo(f"❌ Error: {e}", err=True)
     if verbose:
-        import traceback
         click.echo(traceback.format_exc(), err=True)
     sys.exit(1)
 
@@ -271,8 +272,6 @@ def create_progress_callback(description: str, total: int):
     :return: Callable progress callback function that takes an integer
              representing the number of completed items
     """
-    import sys
-    from doctra.utils.progress import create_beautiful_progress_bar, create_notebook_friendly_bar
 
     # Enhanced environment detection
     is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules

doctra/engines/image_restoration/docres_engine.py

@@ -18,6 +18,8 @@ import sys
 import cv2
 import numpy as np
 import torch
+import tempfile
+import time
 from pathlib import Path
 from typing import Union, List, Tuple, Optional, Dict, Any
 
@@ -85,12 +87,12 @@ def load_docres_weights_from_hf():
         if is_notebook:
             progress_bar = create_notebook_friendly_bar(
                 total=2, 
-                desc="🔄 Downloading DocRes models from Hugging Face Hub"
+                desc="Downloading DocRes models from Hugging Face Hub"
             )
         else:
             progress_bar = create_beautiful_progress_bar(
                 total=2, 
-                desc="🔄 Downloading DocRes models from Hugging Face Hub",
+                desc="Downloading DocRes models from Hugging Face Hub",
                 leave=True
             )
         
@@ -308,8 +310,6 @@ class DocResEngine:
     
     def _run_single_task(self, img_array: np.ndarray, task: str, save_prompts: bool) -> Tuple[np.ndarray, Dict]:
         """Run a single restoration task"""
-        import tempfile
-        import time
         
         # Create temporary file for inference
         with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
@@ -322,7 +322,6 @@ class DocResEngine:
             os.chdir(str(docres_dir))
             
             # Set global DEVICE variable that DocRes inference expects
-            import torch
             import inference  # Import the inference module to set its global DEVICE
             inference.DEVICE = self.device
             
@@ -364,8 +363,6 @@ class DocResEngine:
     
     def _run_end2end_pipeline(self, img_array: np.ndarray, save_prompts: bool) -> Tuple[np.ndarray, Dict]:
         """Run the end2end pipeline: dewarping → deshadowing → appearance"""
-        import tempfile
-        import time
         
         intermediate_steps = {}
         
@@ -374,7 +371,6 @@ class DocResEngine:
         os.chdir(str(docres_dir))
         
         # Set global DEVICE variable that DocRes inference expects
-        import torch
         import inference  # Import the inference module to set its global DEVICE
         inference.DEVICE = self.device
         
@@ -482,7 +478,6 @@ class DocResEngine:
         """
         try:
             from PIL import Image
-            import numpy as np
             from doctra.utils.pdf_io import render_pdf_to_images
             
             # Generate output path if not provided
@@ -510,12 +505,12 @@ class DocResEngine:
             if is_notebook:
                 progress_bar = create_notebook_friendly_bar(
                     total=len(pil_pages), 
-                    desc="🔄 Processing pages"
+                    desc="Processing pages"
                 )
             else:
                 progress_bar = create_beautiful_progress_bar(
                     total=len(pil_pages), 
-                    desc="🔄 Processing pages",
+                    desc="Processing pages",
                     leave=True
                 )
             

doctra/engines/vlm/outlines_types.py

@@ -1,17 +1,19 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 class Chart(BaseModel):
     """
     Structured representation of a chart extracted from an image.
     
-    Contains the title, headers, and data rows extracted from a chart
-    using VLM (Vision Language Model) processing.
+    Includes a title, a short description, column headers, and data rows
+    identified using VLM (Vision Language Model) processing.
 
-    :param title: Title or caption of the chart
+    :param title: Title or caption of the chart (max 31 characters)
+    :param description: Short description of the chart (max 300 characters)
     :param headers: Column headers for the chart data
     :param rows: Data rows containing the chart values
     """
-    title: str
+    title: str = Field(max_length=31)
+    description: str = Field(max_length=300)
     headers: list[str]
     rows: list[list[str]]
 
@@ -19,13 +21,15 @@ class Table(BaseModel):
     """
     Structured representation of a table extracted from an image.
     
-    Contains the title, headers, and data rows extracted from a table
-    using VLM (Vision Language Model) processing.
+    Includes a title, a short description, column headers, and data rows
+    identified using VLM (Vision Language Model) processing.
 
-    :param title: Title or caption of the table
+    :param title: Title or caption of the table (max 31 characters)
+    :param description: Short description of the table (max 300 characters)
     :param headers: Column headers for the table data
     :param rows: Data rows containing the table values
     """
-    title: str
+    title: str = Field(max_length=31)
+    description: str = Field(max_length=300)
     headers: list[str]
     rows: list[list[str]]

doctra/engines/vlm/service.py

@@ -73,7 +73,7 @@ class VLMStructuredExtractor:
         Extract structured chart data from an image.
 
         :param image_path: Path to the chart image file
-        :return: Chart object containing extracted title, headers, and data rows
+        :return: Chart object containing extracted title, description, headers, and data rows
         :raises Exception: If image processing or VLM extraction fails
         """
         prompt_text = (
@@ -81,6 +81,7 @@ class VLMStructuredExtractor:
             "If the title is not present in the image, generate a suitable title. "
             "Ensure that the table represents the data from the chart accurately."
             "The number of columns in the headers must match the number of columns in each row."
+            "Also provide a short description (max 300 characters) of the chart."
         )
         return self._call(prompt_text, image_path, Chart)
 
@@ -89,7 +90,7 @@ class VLMStructuredExtractor:
         Extract structured table data from an image.
 
         :param image_path: Path to the table image file
-        :return: Table object containing extracted title, headers, and data rows
+        :return: Table object containing extracted title, description, headers, and data rows
         :raises Exception: If image processing or VLM extraction fails
         """
         prompt_text = (
@@ -97,5 +98,6 @@ class VLMStructuredExtractor:
             "Provide the headers and rows of the table, ensuring accuracy in the extraction. "
             "If the title is not present in the image, generate a suitable title."
             "The number of columns in the headers must match the number of columns in each row."
+            "Also provide a short description (max 300 characters) of the table."
         )
         return self._call(prompt_text, image_path, Table)

doctra/exporters/excel_writer.py

@@ -5,6 +5,7 @@ from typing import Dict, Any, List, Set
 import pandas as pd  # pip install pandas openpyxl
 from openpyxl.styles import PatternFill, Font, Alignment
 from openpyxl.utils import get_column_letter
+from openpyxl.worksheet.hyperlink import Hyperlink
 
 _INVALID_SHEET_CHARS = r'[:\\/*?\[\]]'  # Excel-invalid characters
 _MAX_SHEET_LEN = 31
@@ -85,6 +86,61 @@ def _autosize_columns(ws, df: pd.DataFrame) -> None:
         ws.column_dimensions[get_column_letter(i)].width = min(max(10, max_len + 2), 60)
 
 
+def _style_summary_sheet(ws, df: pd.DataFrame, sheet_mapping: dict = None) -> None:
+    """
+    Apply special styling to the summary sheet with text wrapping for descriptions.
+    Add hyperlinks to table titles that link to their corresponding sheets.
+    
+    :param ws: OpenPyXL worksheet object to style
+    :param df: Pandas DataFrame containing the summary data
+    :param sheet_mapping: Dictionary mapping table titles to their sheet names
+    :return: None
+    """
+    # Style header row
+    _style_header(ws, ncols=df.shape[1])
+    
+    # Apply text wrapping to all data cells
+    wrap_alignment = Alignment(wrap_text=True, vertical="top")
+    
+    # Apply wrapping to all data rows (skip header row)
+    for row_idx in range(2, len(df) + 2):  # Start from row 2 (after header)
+        for col_idx in range(1, df.shape[1] + 1):
+            cell = ws.cell(row=row_idx, column=col_idx)
+            cell.alignment = wrap_alignment
+            
+            # Add hyperlink to table title column (column A)
+            if col_idx == 1 and sheet_mapping:  # Table Title column
+                table_title = cell.value
+                if table_title and table_title in sheet_mapping:
+                    sheet_name = sheet_mapping[table_title]
+                    
+                    # Create hyperlink to the sheet using proper Excel format
+                    # Escape sheet name if it contains spaces or special characters
+                    if ' ' in sheet_name or any(char in sheet_name for char in ['[', ']', '*', '?', ':', '\\', '/']):
+                        hyperlink_ref = f"#'{sheet_name}'!A1"
+                    else:
+                        hyperlink_ref = f"#{sheet_name}!A1"
+                    
+                    # Use Hyperlink class with proper parameters
+                    cell.hyperlink = Hyperlink(ref=hyperlink_ref, target=hyperlink_ref)
+                    # Style the hyperlink
+                    cell.font = Font(color="0000FF", underline="single")
+    
+    # Set specific column widths for summary sheet
+    # Table Title column - narrower
+    ws.column_dimensions['A'].width = 30
+    # Description column - wider to accommodate wrapped text
+    ws.column_dimensions['B'].width = 60
+    # Page column - narrow for page numbers
+    ws.column_dimensions['C'].width = 10
+    # Type column - narrow for Table/Chart
+    ws.column_dimensions['D'].width = 12
+    
+    # Set row heights to accommodate wrapped text
+    for row_idx in range(2, len(df) + 2):
+        ws.row_dimensions[row_idx].height = 60  # Allow for multiple lines
+
+
 def _normalize_data(headers: List[str], rows: List[List]) -> tuple[List[str], List[List]]:
     """
     Normalize headers and rows to ensure consistent dimensions.
@@ -159,6 +215,31 @@ def write_structured_excel(excel_path: str, items: List[Dict[str, Any]]) -> str
     taken: Set[str] = set()
 
     with pd.ExcelWriter(excel_path, engine="openpyxl", mode="w") as writer:
+        # Create summary sheet first
+        summary_data = []
+        sheet_mapping = {}  # Map table titles to their sheet names
+        
+        for item in valid_items:
+            title = item.get("title") or "Untitled"
+            description = item.get("description") or "No description available"
+            page_number = item.get("page", "Unknown")
+            item_type = item.get("type", "Table")  # Default to "Table" if not specified
+            
+            
+            summary_data.append({
+                "Table Title": title,
+                "Description": description,
+                "Page": page_number,
+                "Type": item_type
+            })
+        
+        # Create summary sheet first (but without hyperlinks initially)
+        if summary_data:
+            summary_df = pd.DataFrame(summary_data)
+            summary_df.to_excel(writer, sheet_name="Table Summary", index=False)
+            taken.add("Table Summary")
+
+        # Process individual table sheets to build sheet mapping
         for item in valid_items:
             try:
                 title = item.get("title") or "Untitled"
@@ -166,6 +247,9 @@ def write_structured_excel(excel_path: str, items: List[Dict[str, Any]]) -> str
                 rows = item.get("rows") or []
 
                 sheet_name = _safe_sheet_name(title, taken)
+                
+                # Add to sheet mapping for hyperlinks
+                sheet_mapping[title] = sheet_name
 
                 # Normalize data to handle mismatched dimensions
                 normalized_headers, normalized_rows = _normalize_data(headers, rows)
@@ -194,4 +278,9 @@ def write_structured_excel(excel_path: str, items: List[Dict[str, Any]]) -> str
                 print(f"Error processing item '{item.get('title', 'Unknown')}': {e}")
                 continue
 
+        # Now add hyperlinks to the summary sheet (after all sheets are created)
+        if summary_data and sheet_mapping:
+            summary_ws = writer.sheets["Table Summary"]
+            _style_summary_sheet(summary_ws, summary_df, sheet_mapping)
+
     return excel_path

doctra/exporters/html_writer.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import os
 import re
 import base64
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 from markdown_it import MarkdownIt
 
 
@@ -64,6 +64,114 @@ def _process_image_paths(md_content: str, out_dir: str) -> str:
     return processed_content
 
 
+def write_html_from_lines(html_lines: List[str], out_dir: str, filename: str = "result.html") -> str:
+    """
+    Convert HTML lines directly into a single HTML file and save it.
+    
+    This function is used when VLM is enabled to ensure proper HTML table formatting
+    instead of markdown-to-HTML conversion.
+
+    :param html_lines: List of HTML strings to join into a single file
+    :param out_dir: Directory where the HTML file will be saved
+    :param filename: Name of the HTML file (default: "result.html")
+    :return: The absolute path of the written HTML file
+    """
+    os.makedirs(out_dir, exist_ok=True)
+
+    # Join HTML lines and clean up excessive blank lines
+    html_content = "\n".join(html_lines).strip() + "\n"
+    html_content = re.sub(r"\n{3,}", "\n\n", html_content)
+
+    # Process image paths to convert relative paths to absolute paths or base64
+    html_content = _process_image_paths(html_content, out_dir)
+    
+    # Always apply table styling to ensure all tables are properly formatted
+    html_content = _add_table_styling(html_content)
+    
+    # Create complete HTML document with modern styling
+    html_document = f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Document Analysis Results</title>
+    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap" rel="stylesheet">
+    <style>
+        {_get_css_styles()}
+    </style>
+</head>
+<body>
+    <button class="theme-toggle" onclick="toggleTheme()" title="Toggle dark mode"></button>
+    <div class="container">
+        <header class="header">
+            <div class="header-content">
+                <div class="header-text">
+                    <h1>Document Analysis Results</h1>
+                    <p class="subtitle">Intelligent Document Processing & Analysis</p>
+                </div>
+                <div class="header-badge">
+                    Generated by Doctra
+                </div>
+            </div>
+        </header>
+        <main class="content">
+            {html_content}
+        </main>
+        <footer class="footer">
+            <div class="footer-content">
+                <div class="footer-brand">Doctra</div>
+                <div class="footer-info">
+                    <span>Intelligent Document Processing</span>
+                    <a href="https://github.com/AdemBoukhris457/Doctra" target="_blank">GitHub</a>
+                </div>
+            </div>
+        </footer>
+    </div>
+    <script>
+        // Theme toggle functionality
+        function toggleTheme() {{
+            const body = document.body;
+            const currentTheme = body.getAttribute('data-theme');
+            const newTheme = currentTheme === 'dark' ? 'light' : 'dark';
+            
+            body.setAttribute('data-theme', newTheme);
+            localStorage.setItem('doctra-theme', newTheme);
+            
+            // Add smooth transition
+            body.style.transition = 'all 0.3s ease';
+            setTimeout(() => {{
+                body.style.transition = '';
+            }}, 300);
+        }}
+
+        // Load saved theme on page load
+        document.addEventListener('DOMContentLoaded', function() {{
+            const savedTheme = localStorage.getItem('doctra-theme') || 'light';
+            document.body.setAttribute('data-theme', savedTheme);
+        }});
+
+        // Add smooth scroll behavior
+        document.documentElement.style.scrollBehavior = 'smooth';
+
+        // Add loading animation
+        window.addEventListener('load', function() {{
+            document.body.style.opacity = '0';
+            document.body.style.transition = 'opacity 0.5s ease';
+            setTimeout(() => {{
+                document.body.style.opacity = '1';
+            }}, 100);
+        }});
+    </script>
+</body>
+</html>"""
+
+    html_path = os.path.join(out_dir, filename)
+    with open(html_path, "w", encoding="utf-8") as f:
+        f.write(html_document)
+
+    return os.path.abspath(html_path)
+
+
 def write_html(md_lines: List[str], out_dir: str, filename: str = "result.html") -> str:
     """
     Convert collected Markdown lines into a single HTML file and save it.
@@ -414,6 +522,54 @@ def _create_html_table(headers: List[str], rows: List[List]) -> str:
     """
 
 
+def render_html_table(
+    headers: List[str] | None,
+    rows: List[List[str]] | None,
+    title: Optional[str] = None,
+) -> str:
+    """
+    Render an HTML table from headers, rows, and optional title.
+    
+    Creates a properly formatted HTML table with headers, data rows,
+    and optional title. This is used for VLM-extracted tables to ensure
+    they display as proper HTML tables instead of markdown.
+
+    :param headers: List of column headers (optional, will be auto-generated if None)
+    :param rows: List of data rows, where each row is a list of cell values
+    :param title: Optional title to display above the table
+    :return: Formatted HTML table string
+    """
+    headers = headers or []
+    rows = rows or []
+
+    if not headers and not rows:
+        return "<p class='no-data'>No data available</p>"
+
+    # Determine width
+    width = len(headers) if headers else (max((len(r) for r in rows), default=1))
+
+    # Generate headers if not provided
+    if not headers:
+        headers = [f"Column {i+1}" for i in range(width)]
+
+    # Normalize data to handle mismatched dimensions
+    normalized_headers, normalized_rows = _normalize_data(headers, rows)
+
+    # Create HTML table
+    table_html = _create_html_table(normalized_headers, normalized_rows)
+    
+    # Add title if provided
+    if title:
+        return f"""
+        <div class="table-section">
+            <h3 class="table-title">{_escape_html(title)}</h3>
+            {table_html}
+        </div>
+        """
+    else:
+        return table_html
+
+
 def _add_table_styling(html_content: str) -> str:
     """
     Add table styling wrapper to HTML content.
@@ -884,6 +1040,55 @@ def _get_css_styles() -> str:
             content: '☀️';
         }
 
+        /* Dark mode table styles */
+        [data-theme="dark"] .markdown-table,
+        [data-theme="dark"] table {
+            background: var(--card-bg);
+            border-color: var(--border-color);
+        }
+
+        [data-theme="dark"] .markdown-table th,
+        [data-theme="dark"] table th {
+            background: #374151;
+            color: #f9fafb;
+            border-bottom-color: var(--accent-color);
+        }
+
+        [data-theme="dark"] .markdown-table td,
+        [data-theme="dark"] table td {
+            color: #f9fafb;
+            border-bottom-color: var(--border-color);
+        }
+
+        [data-theme="dark"] .markdown-table tr:nth-child(even),
+        [data-theme="dark"] table tr:nth-child(even) {
+            background: #374151;
+        }
+
+        [data-theme="dark"] .markdown-table tr:hover,
+        [data-theme="dark"] table tr:hover {
+            background: #4b5563;
+        }
+
+        /* Dark mode footer styles to match header */
+        [data-theme="dark"] .footer {
+            background: var(--primary-color);
+            color: white;
+            border-top-color: var(--accent-color);
+        }
+
+        [data-theme="dark"] .footer-brand {
+            color: white;
+        }
+
+        [data-theme="dark"] .footer a {
+            color: rgba(255, 255, 255, 0.8);
+        }
+
+        [data-theme="dark"] .footer a:hover {
+            color: white;
+        }
+
         /* Professional scrollbar */
         ::-webkit-scrollbar {
             width: 8px;